package com.xizi.RBTree;

/**
 * 手撕红黑树
 * 红黑树是一颗二叉查找树
 * 根节点是黑色
 * 每个叶子节点是黑色
 * 如果一个节点是红色 那么它的子节点必须是黑色
 * 从帮一个节点到该节点的子孙节点的所有路径上包含相同数目的黑色节点
 */
public class RBTree {
    RBTreeNode root;
    private final boolean RED = false;
    private final boolean BLACK = true;

    //查询节点
    public RBTreeNode query(int key) {
        RBTreeNode tmp = root;
        //利用二叉查找树的特点进行递归查找 左<中<右
        while (tmp != null) {
            if (tmp.getKey() == key)
                return tmp;
            else if (tmp.getKey() > key)
                tmp = tmp.getLeft();
            else
                tmp = tmp.getRight();
        }
        return null;
    }

    //插入节点
    public void insert(int key) {
        //默认是红色
        RBTreeNode node = new RBTreeNode(key);
        //如果根节点是空 直接插入当前节点作为根节点
        if (root == null) {
            root = node;
            node.setColor(BLACK);
            return;
        }

        //根据当前节点的值 进行判断在树的左边还是右边
        RBTreeNode parent = root;
        RBTreeNode son = null;
        if (key <= parent.getKey()) {
            son = parent.getLeft();
        } else {
            son = parent.getRight();
        }
        //找到当前节点在二叉搜索树中的位置
        while (son != null) {
            parent = son;
            if (key <= parent.getKey()) {
                son = parent.getLeft();
            } else {
                son = parent.getRight();
            }
        }
        //将当前节点插入到树中
        if (key <= parent.getKey()) {
            parent.setLeft(node);
        } else {
            parent.setRight(node);
        }
        node.setParent(parent);

        //进行调整 满足的红黑树性质
        insertFix(node);
    }

    private void insertFix(RBTreeNode node) {
        RBTreeNode father, grandFather;
        while ((father = node.getParent()) != null && father.getColor() == RED) {
            grandFather = father.getParent();
            //第一种情况 F为红，U为红 插入一个新的节点也是红色 在F的左边
            if (grandFather.getLeft() == father) {  //F为G左儿子的情况，如之前的分析
                //U做为新插入节点的叔叔 也是红色
                RBTreeNode uncle = grandFather.getRight();
                if (uncle != null && uncle.getColor() == RED) {
                    //将插入节点的父亲和叔叔的颜色设置为黑色
                    setBlack(father);
                    setBlack(uncle);
                    //将祖父设置为红色
                    setRed(grandFather);
                    //再祖父节点基础上继续查找
                    node = grandFather;
                    continue;
                }
                //如果 叔叔节点为黑色 当前节点作为右节点 父节点进行左旋
                if (node == father.getRight()) {
                    leftRotate(father);
                    //左旋之后进行父节点和当前节点的交换 为了后面的操作
                    RBTreeNode tmp = node;
                    node = father;
                    father = tmp;
                }
                //左旋之后 将父节点设置为黑色
                setBlack(father);
                //将祖父节点设置红色
                setRed(grandFather);
                //在进行祖父节点右旋之后符合红黑树的性质
                rightRotate(grandFather);
            } else {    // F为G的右儿子的情况，对称操作
                RBTreeNode uncle = grandFather.getLeft();
                if (uncle != null && uncle.getColor() == RED) {
                    setBlack(father);
                    setBlack(uncle);
                    setRed(grandFather);
                    node = grandFather;
                    continue;
                }
                if (node == father.getLeft()) {
                    rightRotate(father);
                    RBTreeNode tmp = node;
                    node = father;
                    father = tmp;
                }
                setBlack(father);
                setRed(grandFather);
                leftRotate(grandFather);
            }
        }
        setBlack(root);
    }

    public void delete(int key) {
        delete(query(key));
    }

    private void delete(RBTreeNode node) {
        if (node == null)
            return;
        //当删除的左右子节点都不为空的时候
        if (node.getLeft() != null && node.getRight() != null) {
            //替代节点（选择方式是X的右子树的最小点）
            RBTreeNode replaceNode = node;
            //得到右子树
            RBTreeNode tmp = node.getRight();
            while (tmp != null) {
                replaceNode = tmp;
                //然后向左子树找最小的节点（二叉搜索树）
                tmp = tmp.getLeft();
            }
            //获取到替代节点的值
            int t = replaceNode.getKey();
            replaceNode.setKey(node.getKey());
            //将删除节点的值进行替换
            node.setKey(t);
            //再删除替代节点（已经使用了）
            delete(replaceNode);
            return;
        }
        // 当删除的左右子节点不满足都不为空的时候
        RBTreeNode replaceNode = null;
        if (node.getLeft() != null)
            replaceNode = node.getLeft();
        else
            replaceNode = node.getRight();

        RBTreeNode parent = node.getParent();
        if (parent == null) {
            root = replaceNode;
            if (replaceNode != null)
                replaceNode.setParent(null);
        } else {
            if (replaceNode != null)
                replaceNode.setParent(parent);
            if (parent.getLeft() == node)
                parent.setLeft(replaceNode);
            else {
                parent.setRight(replaceNode);
            }
        }
        if (node.getColor() == BLACK)
            removeFix(parent, replaceNode);

    }

    //多余的颜色在node里
    private void removeFix(RBTreeNode father, RBTreeNode node) {
        while ((node == null || node.getColor() == BLACK) && node != root) {
            if (father.getLeft() == node) {  //S为P的左儿子的情况，如之前的分析
                RBTreeNode brother = father.getRight();
                if (brother != null && brother.getColor() == RED) {
                    setRed(father);
                    setBlack(brother);
                    leftRotate(father);
                    brother = father.getRight();
                }
                if (brother == null || (isBlack(brother.getLeft()) && isBlack(brother.getRight()))) {
                    setRed(brother);
                    node = father;
                    father = node.getParent();
                    continue;
                }
                if (isRed(brother.getLeft())) {
                    setBlack(brother.getLeft());
                    setRed(brother);
                    rightRotate(brother);
                    brother = brother.getParent();
                }

                brother.setColor(father.getColor());
                setBlack(father);
                setBlack(brother.getRight());
                leftRotate(father);
                node = root;//跳出循环
            } else {                         //S为P的右儿子的情况，对称操作
                RBTreeNode brother = father.getLeft();
                if (brother != null && brother.getColor() == RED) {
                    setRed(father);
                    setBlack(brother);
                    rightRotate(father);
                    brother = father.getLeft();
                }
                if (brother == null || (isBlack(brother.getLeft()) && isBlack(brother.getRight()))) {
                    setRed(brother);
                    node = father;
                    father = node.getParent();
                    continue;
                }
                if (isRed(brother.getRight())) {
                    setBlack(brother.getRight());
                    setRed(brother);
                    leftRotate(brother);
                    brother = brother.getParent();
                }

                brother.setColor(father.getColor());
                setBlack(father);
                setBlack(brother.getLeft());
                rightRotate(father);
                node = root;//跳出循环
            }
        }

        if (node != null)
            node.setColor(BLACK);
    }

    private boolean isBlack(RBTreeNode node) {
        if (node == null)
            return true;
        return node.getColor() == BLACK;
    }

    private boolean isRed(RBTreeNode node) {
        if (node == null)
            return false;
        return node.getColor() == RED;
    }

    private void leftRotate(RBTreeNode node) {
        RBTreeNode right = node.getRight();
        RBTreeNode parent = node.getParent();
        if (parent == null) {
            root = right;
            right.setParent(null);
        } else {
            if (parent.getLeft() != null && parent.getLeft() == node) {
                parent.setLeft(right);
            } else {
                parent.setRight(right);
            }
            right.setParent(parent);
        }
        node.setParent(right);
        node.setRight(right.getLeft());
        if (right.getLeft() != null) {
            right.getLeft().setParent(node);
        }
        right.setLeft(node);
    }

    private void rightRotate(RBTreeNode node) {
        RBTreeNode left = node.getLeft();
        RBTreeNode parent = node.getParent();
        if (parent == null) {
            root = left;
            left.setParent(null);
        } else {
            if (parent.getLeft() != null && parent.getLeft() == node) {
                parent.setLeft(left);
            } else {
                parent.setRight(left);
            }
            left.setParent(parent);
        }
        node.setParent(left);
        node.setLeft(left.getRight());
        if (left.getRight() != null) {
            left.getRight().setParent(node);
        }
        left.setRight(node);
    }

    private void setBlack(RBTreeNode node) {
        node.setColor(BLACK);
    }

    private void setRed(RBTreeNode node) {
        node.setColor(RED);
    }

    public void inOrder() {
        inOrder(root);
    }

    private void inOrder(RBTreeNode node) {
        if (node == null)
            return;
        inOrder(node.getLeft());
        System.out.println(node);
        inOrder(node.getRight());
    }
}

